import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
from sklearn.neighbors import KNeighborsClassifier

# 创建训练数据集
def getData():
    x_train = np.array([[3.393533211, 2.331273381],
                        [3.110073483, 1.781539638],
                        [1.343808831, 3.368360954],
                        [3.582294042, 4.679179110],
                        [2.280362439, 2.866990263],
                        [7.423436942, 4.696522875],
                        [5.745051997, 3.533989803],
                        [9.172168622, 2.511101045],
                        [7.792783481, 3.424088941],
                        [7.939820817, 0.791637231]
                        ])
    y_train = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
    return x_train,y_train

# kNN算法
def kNN(k,x_train,y_train,pre_data):
    for x in pre_data:
        dis=[np.sqrt(np.sum((xi_train-x)**2)) for xi_train in x_train]
        # argsort得到的是原数组的索引
        nearest=np.argsort(dis)
        topK_y=[y_train[neighbor] for neighbor in nearest[:k]]
        votes=Counter(topK_y).most_common()
        print(str(pre_data[0])+"---> predict type --->"+str(votes[0][0]))
        plt.scatter(x[0],x[1],color='b')

    plt.scatter(x_train[y_train==0,0],x_train[y_train==0,1],color='g')
    plt.scatter(x_train[y_train==1,0],x_train[y_train==1,1],color='r')
    plt.show()

if __name__=='__main__':
    x_train,y_train=getData()
    x_test=[
        [8.093607318, 3.365731514],
        [7.254879955, 8.745666941]
    ]
    k=6
    kNN(k,x_train,y_train,x_test)

    # 使用scikit-learn中的kNN
    kNN_classifier=KNeighborsClassifier(n_neighbors=6)
    kNN_classifier.fit(x_train,y_train)
    res=kNN_classifier.predict(x_test)
    print(res)